#include "crypt.h"
#include <unistd.h>
#include <string.h>
#include <fcntl.h>
#include <stdlib.h>
#include <memory.h>

#define BUFSIZE 4096

void usage() {
  printf("usage\n");
  printf("生成公私钥:rsa --genkey -p passwd\n");
  printf("公钥加密:rsa -e file -k public.key -o outfile\n");
  printf("私钥解密:rsa -d file -k private.key -o putfile\n");
}

int main(int argc,char **argv) {

  //生成密钥 rsa --genkey -p password
  if(argc==4) {
    if(strcmp(argv[1],"--genkey")==0&&strcmp(argv[2],"-p")==0) {
      if(rsa_generate_key_files("public.key", "private.key",argv[3])<0) {
        printf("failed to generate key\n");
      } else{
        printf("public.key private.key success\n");
      }
      return 0;
    }
  }

  //加解密
  if(argc==7) {
    if((strcmp(argv[1],"-e")==0||strcmp(argv[1],"-d")==0)&&strcmp(argv[3],"-k")==0&&strcmp(argv[5],"-o")==0) {
      //打开文件
      int r_fd = open(argv[2], O_RDONLY);
      int w_fd = open(argv[6], O_WRONLY | O_CREAT | O_TRUNC,00600);
      if(r_fd==-1||w_fd==-1) {
        printf("failed to open file\n");
        return 0;
      }
      char read_buf[BUFSIZE+256]={0},write_buf[BUFSIZE+256]={0};
      char ivec[]=IVEC;
      //加密
      if(strcmp(argv[1],"-e")==0) {
        //生成aes密钥
        char aes_key[AES_KEY_LEN*2];
        random_generate_aeskey(aes_key,AES_KEY_LEN/8);
        //取出公钥
        char pub_key[2048];
        read_key_file(argv[4],pub_key);
        //使用rsa把aes_key加密，并放在out文件的前512字节处
        rsa_pub_encrypt(aes_key,AES_KEY_LEN,write_buf,pub_key);
        write(w_fd,write_buf,512);
        //读取明文文件，使用aes加密后写入密文文件中
        int len;
        while((len = read(r_fd,read_buf,BUFSIZE))>0) {
          if(len<BUFSIZE) {
            //使用PKCS7填充
          	size_t padd_num = 16 - len % 16;
          	for(int i=0;i<padd_num;i++)
              read_buf[len++]=padd_num;
          }
          len = aes_encrypt(read_buf,len,write_buf,aes_key,ivec);
          write(w_fd,write_buf,len);
        }
      } else { //解密
        //先取出密文的前512字节
        read(r_fd,read_buf,512);
        //取出私钥
        char priv_key[2048];
        read_key_file(argv[4],priv_key);
        //解密得到aes_key
        char aes_key[AES_KEY_LEN*2];
        rsa_priv_decrypt(read_buf,512,aes_key,priv_key);
        //通过aes_key解密密文文件并写到目标文件中
        int len,temp;//temp保存上次加密后write_buf中密文的长度，用于最后去填充
        while ((len=read(r_fd,read_buf,BUFSIZE))>0) {
          aes_decrypt(read_buf,len,write_buf,aes_key,ivec);
          write(w_fd,write_buf,len);
          temp=len;
        }
        //去填充
        len = lseek(w_fd,0,SEEK_END);
        ftruncate(w_fd,len-(int)(write_buf[temp-1]));
      }
      close(r_fd);close(w_fd);
      return 0;
    }
  }
  usage();
  return 0;
}
